#include <iostream>
#include <openssl/bn.h>
#include "nprime_and_rsquare.hpp"

#ifndef BYTELEN
#define BYTELEN 8
#endif
#ifndef BYTECOUNT
#define BYTECOUNT(bitcount) ((bitcount + BYTELEN - 1) / BYTELEN)
#endif
#ifndef MON_MUL_BLOCK_LENGTH
#define MON_MUL_BLOCK_LENGTH 128
#endif

void gen_n_prime_and_r_square(char *p, int p_bitcount, char *n_prime, char *r_square) {
  BIGNUM *bn_p = BN_new();
  BIGNUM *temp = BN_new();
  BN_bin2bn((unsigned char*)p, BYTECOUNT(p_bitcount), bn_p);
//  std::cout << "p(dec): " << BN_bn2dec(bn_p) << std::endl;
//  std::cout << "p(hex): " << BN_bn2hex(bn_p) << std::endl;
  BIGNUM *bn_n_prime = BN_new();
  BIGNUM *bn_r_square = BN_new();
  BN_CTX *bn_ctx = BN_CTX_new();
  // compute n_prime = -p^(-1) mod 2^block_bits
  BN_lshift(bn_r_square, BN_value_one(), MON_MUL_BLOCK_LENGTH);
  BN_mod_inverse(bn_n_prime, bn_p, bn_r_square, bn_ctx);
  BN_zero(temp);
  BN_sub(bn_n_prime, temp, bn_n_prime);
  BN_nnmod(bn_n_prime, bn_n_prime, bn_r_square, bn_ctx);
  // compute 2^(2 * (n + block_bits)) mod p
  BN_lshift(bn_r_square, BN_value_one(), 2 * (p_bitcount + MON_MUL_BLOCK_LENGTH));
  BN_mod(bn_r_square, bn_r_square, bn_p, bn_ctx);
  BN_bn2binpad(bn_n_prime, (unsigned char*)(n_prime), BYTECOUNT(MON_MUL_BLOCK_LENGTH));
  BN_bn2binpad(bn_r_square, (unsigned char*)(r_square), BYTECOUNT(p_bitcount));
//  std::cout << "n_prime(dec): " << BN_bn2dec(bn_n_prime) << std::endl;
//  std::cout << "n_prime(hex): " << BN_bn2hex(bn_n_prime) << std::endl;
//  std::cout << "r_square(dec): " << BN_bn2dec(bn_r_square) << std::endl;
//  std::cout << "r_square(hex): " << BN_bn2hex(bn_r_square) << std::endl;
}
